import os
import numpy as np
import json
from openai import OpenAI
import fitz
from tqdm import tqdm

# 设置 API 变量
os.environ["OPENAI_API_KEY"] = "sk-aswwzdvvimeybiiokqebixpkhbmcftlbgkubssfuodifqjcf"
# 初始化OpenAI客户端
client = OpenAI(
    base_url="https://api.siliconflow.cn/v1",
    api_key=os.getenv("OPENAI_API_KEY")  # 获取环境变量
)

def create_embeddings(text, model="BAAI/bge-m3"):
    response = client.embeddings.create(
        model=model,
        input=text
    )
    return response.data[0].embedding

def extract_text_from_pdf(pdf_path):
    mypdf = fitz.open(pdf_path)
    all_text = ""
    for page_num in range(mypdf.page_count):
        page = mypdf[page_num]
        text = page.get_text("text")
        all_text += text
    return all_text

def generate_chunk_header(chunk, model="Qwen/Qwen2-1.5B-Instruct"):
    system_prompt = "Generate a concise and informative title for the given text."
    response = client.chat.completions.create(
        model=model,
        temperature=0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": chunk}
        ]
    )
    return response.choices[0].message.content.strip()

def chunk_text_with_headers(text, n, overlap):
    chunks = []
    for i in range(0, len(text), n - overlap):
        chunk = text[i:i + n]
        header = generate_chunk_header(chunk)
        chunks.append({"header": header, "text": chunk})
    return chunks  # Return the list of chunks with headers

def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

def semantic_search(query, chunks, k=2):
    query_embedding = create_embeddings(query)
    similarities = []
    for chunk in chunks:
        sim_text = cosine_similarity(np.array(query_embedding), np.array(chunk["embedding"]))
        sim_header = cosine_similarity(np.array(query_embedding), np.array(chunk["header_embedding"]))
        avg_similarity = (sim_text + sim_header) / 2
        similarities.append((chunk, avg_similarity))
    similarities.sort(key=lambda x: x[1], reverse=True)
    return [x[0] for x in similarities[:k]]

def generate_response(system_prompt, user_message, model="Qwen/Qwen2-1.5B-Instruct"):
    response = client.chat.completions.create(
        model=model,
        temperature=0,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    )
    return response

pdf_path = "data/AI_Information.pdf"
extracted_text = extract_text_from_pdf(pdf_path)
text_chunks = chunk_text_with_headers(extracted_text, 1000, 200)
print("Sample Chunk:")
print("Header:", text_chunks[0]['header'])
print("Content:", text_chunks[0]['text'])
embeddings = []  # Initialize an empty list to store embeddings
for chunk in tqdm(text_chunks, desc="Generating embeddings"):
    text_embedding = create_embeddings(chunk["text"])
    header_embedding = create_embeddings(chunk["header"])
    embeddings.append({"header": chunk["header"], "text": chunk["text"], "embedding": text_embedding, "header_embedding": header_embedding})
with open('data/val.json') as f:
    data = json.load(f)
query = data[0]['question']
top_chunks = semantic_search(query, embeddings, k=2)
print("Query:", query,"\n")
for i, chunk in enumerate(top_chunks):
    print(f"Header {i+1}: {chunk['header']}")
    print(f"Content:{chunk['text']}\n")
system_prompt = ("You are an AI assistant that strictly answers based on the given context. "
                 "If the answer cannot be derived directly from the provided context, "
                 "respond with: 'I do not have enough information to answer that.'")
user_prompt = "\n".join([f"Header: {chunk['header']}\nContent:\n{chunk['text']}" for chunk in top_chunks])
user_prompt = f"{user_prompt}\nQuestion: {query}"
ai_response = generate_response(system_prompt, user_prompt)
evaluate_system_prompt = """You are an intelligent evaluation system. 
Assess the AI assistant's response based on the provided context. 
- Assign a score of 1 if the response is very close to the true answer. 
- Assign a score of 0.00~1.00 if the response is partially correct. 
- Assign a score of 0 if the response is incorrect."""
true_answer = data[0]['ideal_answer']
evaluation_prompt = f"""
User Query: {query}
AI Response: {ai_response}
True Answer: {true_answer}
{evaluate_system_prompt}
"""
evaluation_response = generate_response(evaluate_system_prompt, evaluation_prompt)
print("Evaluation Score:", evaluation_response.choices[0].message.content)